import torch
import numpy as np

# 获取最优策略
def get_optimal_policy(net, states):
    q_table = torch.zeros((25, 5))
    with torch.no_grad():
        for i in range(5):
            for j in range(5):
                q_table[i * 5 + j] = net(states[i * 5 + j])
    policy = q_table.argmax(dim=1).reshape((5, 5))
    return policy


# 可视化策略
def visual_policy(policy,show=True):
    signs = np.array(["\u2191", "\u2193", "\u2190", "\u2192", "\u25CB"])
    graph = np.zeros((5, 5), dtype=str)
    for i in range(5):
        for j in range(5):
            graph[i][j] = signs[policy[i][j]]
    if show:
        print(graph)
    return graph.__str__()